import numpy as np
import matplotlib.pyplot as plt
import os
import time

from vehicle_model import DoubleIntegratorVehicle
from lqr_controller import LQRController
from trajectory_generator import TrajectoryGenerator
from simulation import Simulation
from visualization import Visualization
from utils import set_plot_style

def test_trajectory_generation():
    """测试轨迹生成"""
    print("\n=== 测试轨迹生成 ===")
    
    dt = 0.1
    sim_time = 30
    time_points = np.arange(0, sim_time, dt)
    
    # 创建轨迹生成器
    traj_gen = TrajectoryGenerator(dt=dt)
    
    # 生成不同类型的轨迹
    traj_circle = traj_gen.generate_trajectory('circle', time_points)
    traj_figure_eight = traj_gen.generate_trajectory('figure_eight', time_points)
    traj_square = traj_gen.generate_trajectory('square', time_points)
    
    # 绘制轨迹
    plt.figure(figsize=(15, 5))
    
    plt.subplot(1, 3, 1)
    plt.plot(traj_circle[:, 0], traj_circle[:, 1], 'b-')
    plt.title('圆形轨迹')
    plt.grid(True)
    plt.axis('equal')
    
    plt.subplot(1, 3, 2)
    plt.plot(traj_figure_eight[:, 0], traj_figure_eight[:, 1], 'g-')
    plt.title('8字形轨迹')
    plt.grid(True)
    plt.axis('equal')
    
    plt.subplot(1, 3, 3)
    plt.plot(traj_square[:, 0], traj_square[:, 1], 'r-')
    plt.title('方形轨迹')
    plt.grid(True)
    plt.axis('equal')
    
    plt.tight_layout()
    plt.show()
    
    print("轨迹生成测试完成!")
    return traj_circle, traj_figure_eight, traj_square

def test_lqr_controller():
    """测试LQR控制器"""
    print("\n=== 测试LQR控制器 ===")
    
    # 创建双积分小车模型
    dt = 0.1
    vehicle = DoubleIntegratorVehicle(dt=dt)
    
    # 获取系统矩阵
    A, B = vehicle.get_state_space_matrices()
    
    print("系统矩阵:")
    print("A =\n", A)
    print("B =\n", B)
    
    # 测试不同的Q和R值
    Q_values = [
        np.diag([10, 10, 1, 1]),     # 基础
        np.diag([100, 100, 10, 10]), # 更高的位置和速度权重
        np.diag([500, 500, 50, 50])  # 非常高的权重
    ]
    
    R_values = [
        np.diag([1, 1]),      # 基础
        np.diag([0.1, 0.1]),  # 低控制权重
        np.diag([0.01, 0.01]) # 非常低的控制权重
    ]
    
    print("\n测试不同的Q和R值:")
    
    for i, Q in enumerate(Q_values):
        for j, R in enumerate(R_values):
            print(f"\nQ{i+1} = \n{Q}")
            print(f"R{j+1} = \n{R}")
            
            # 创建LQR控制器
            controller = LQRController(A, B, Q, R)
            
            # 测试控制输入计算
            current_state = np.array([0, 0, 0, 0])
            reference_state = np.array([1, 1, 0, 0])
            
            u = controller.compute_control_input(current_state, reference_state)
            print(f"控制输入 u = {u}")
    
    print("\nLQR控制器测试完成!")

def verify_system_update():
    """验证系统状态更新计算"""
    print("\n=== 验证系统状态更新 ===")
    
    # 创建双积分小车模型
    dt = 0.1
    vehicle = DoubleIntegratorVehicle(dt=dt)
    
    # 设置初始状态
    initial_state = np.array([0, 0, 0, 0])
    vehicle.reset(initial_state)
    
    # 应用固定控制输入并观察状态变化
    control_input = np.array([1, 1])  # 在x和y方向都施加单位加速度
    
    print("初始状态:", vehicle.get_state())
    
    # 跟踪状态历史
    states = [vehicle.get_state()]
    
    # 运行10个时间步
    for i in range(10):
        new_state = vehicle.step(control_input)
        states.append(new_state)
        print(f"时间步 {i+1}, 控制输入 {control_input}, 新状态 {new_state}")
    
    # 将状态历史转换为numpy数组
    states = np.array(states)
    
    # 绘制位置随时间变化图
    plt.figure(figsize=(12, 8))
    
    plt.subplot(2, 2, 1)
    plt.plot(states[:, 0], states[:, 1], 'bo-')
    plt.xlabel('X位置')
    plt.ylabel('Y位置')
    plt.title('位置轨迹')
    plt.grid(True)
    
    plt.subplot(2, 2, 2)
    time_points = np.arange(0, 11) * dt
    plt.plot(time_points, states[:, 0], 'r-', label='X')
    plt.plot(time_points, states[:, 1], 'g-', label='Y')
    plt.xlabel('时间(s)')
    plt.ylabel('位置')
    plt.title('位置随时间变化')
    plt.legend()
    plt.grid(True)
    
    plt.subplot(2, 2, 3)
    plt.plot(time_points, states[:, 2], 'r-', label='Vx')
    plt.plot(time_points, states[:, 3], 'g-', label='Vy')
    plt.xlabel('时间(s)')
    plt.ylabel('速度')
    plt.title('速度随时间变化')
    plt.legend()
    plt.grid(True)
    
    plt.tight_layout()
    plt.show()
    
    print("系统状态更新验证完成!")

def debug_lqr_tracking():
    """调试LQR轨迹跟踪"""
    print("\n=== 调试LQR轨迹跟踪 ===")
    
    # 参数设置
    dt = 0.1
    sim_time = 30
    
    # 创建双积分小车模型
    vehicle = DoubleIntegratorVehicle(dt=dt)
    
    # 获取系统矩阵
    A, B = vehicle.get_state_space_matrices()
    
    # 创建轨迹生成器和生成轨迹
    traj_gen = TrajectoryGenerator(dt=dt)
    time_points = np.arange(0, sim_time, dt)
    trajectory = traj_gen.generate_trajectory('circle', time_points)
    
    # 不同的Q和R配置
    configs = [
        {"name": "基础", "Q": np.diag([10, 10, 1, 1]), "R": np.diag([1, 1])},
        {"name": "高位置权重", "Q": np.diag([100, 100, 1, 1]), "R": np.diag([1, 1])},
        {"name": "低控制权重", "Q": np.diag([10, 10, 1, 1]), "R": np.diag([0.1, 0.1])},
        {"name": "最优配置", "Q": np.diag([500, 500, 50, 50]), "R": np.diag([0.01, 0.01])}
    ]
    
    # 存储不同配置的结果
    all_results = []
    
    # 对每个配置运行仿真
    for config in configs:
        print(f"\n测试配置: {config['name']}")
        print(f"Q = \n{config['Q']}")
        print(f"R = \n{config['R']}")
        
        # 创建LQR控制器
        controller = LQRController(A, B, config['Q'], config['R'])
        
        # 设置初始状态
        initial_state = trajectory[0].copy()
        current_state = initial_state.copy()
        vehicle.reset(current_state)
        
        # 跟踪状态历史
        states = [current_state.copy()]
        controls = []
        
        # 运行仿真
        for i in range(len(time_points) - 1):
            # 获取参考状态
            reference_state = trajectory[i]
            
            # 计算控制输入
            u = controller.compute_control_input(current_state, reference_state)
            controls.append(u)
            
            # 更新状态
            current_state = vehicle.step(u)
            states.append(current_state.copy())
            
            # 打印调试信息（仅前几步）
            if i < 5:
                print(f"  Step {i}:")
                print(f"    参考状态: {reference_state}")
                print(f"    当前状态: {current_state}")
                print(f"    控制输入: {u}")
                print(f"    位置误差: {np.linalg.norm(reference_state[:2] - current_state[:2])}")
        
        # 转换为numpy数组
        states = np.array(states)
        controls = np.array(controls)
        
        # 计算跟踪误差
        position_errors = trajectory[:len(states), :2] - states[:, :2]
        distance_errors = np.sqrt(np.sum(position_errors**2, axis=1))
        rmse = np.sqrt(np.mean(distance_errors**2))
        
        print(f"  RMSE: {rmse:.4f}")
        
        # 存储结果
        all_results.append({
            "name": config['name'],
            "states": states,
            "controls": controls,
            "rmse": rmse
        })
    
    # 可视化比较
    plt.figure(figsize=(12, 8))
    
    # 绘制轨迹
    plt.subplot(1, 2, 1)
    plt.plot(trajectory[:, 0], trajectory[:, 1], 'k-', label='参考轨迹')
    
    for result in all_results:
        plt.plot(result['states'][:, 0], result['states'][:, 1], '--', label=f"{result['name']} (RMSE: {result['rmse']:.4f})")
    
    plt.xlabel('X位置')
    plt.ylabel('Y位置')
    plt.title('轨迹比较')
    plt.legend()
    plt.grid(True)
    plt.axis('equal')
    
    # 绘制误差
    plt.subplot(1, 2, 2)
    
    for result in all_results:
        position_errors = trajectory[:len(result['states']), :2] - result['states'][:, :2]
        distance_errors = np.sqrt(np.sum(position_errors**2, axis=1))
        plt.plot(time_points[:len(distance_errors)], distance_errors, label=result['name'])
    
    plt.xlabel('时间(s)')
    plt.ylabel('跟踪误差')
    plt.title('跟踪误差随时间变化')
    plt.legend()
    plt.grid(True)
    
    plt.tight_layout()
    plt.show()
    
    print("LQR轨迹跟踪调试完成!")
    
    # 返回最优配置
    best_result = min(all_results, key=lambda x: x['rmse'])
    return best_result

if __name__ == "__main__":
    # 设置绘图样式
    set_plot_style()
    
    # 运行测试
    print("开始诊断与调试程序...")
    
    # 测试轨迹生成
    traj_circle, traj_figure_eight, traj_square = test_trajectory_generation()
    
    # 测试LQR控制器
    test_lqr_controller()
    
    # 验证系统状态更新
    verify_system_update()
    
    # 调试LQR轨迹跟踪
    best_result = debug_lqr_tracking()
    
    print(f"\n调试完成！推荐使用配置: {best_result['name']}")
    print(f"对应的RMSE: {best_result['rmse']:.4f}")
    
    print("\n您可以使用以下命令运行主程序:")
    if best_result['name'] == "最优配置":
        print("python main.py --q-pos 500 --q-vel 50 --r-acc 0.01")
    elif best_result['name'] == "高位置权重":
        print("python main.py --q-pos 100 --q-vel 1 --r-acc 1")
    elif best_result['name'] == "低控制权重":
        print("python main.py --q-pos 10 --q-vel 1 --r-acc 0.1")
    else:
        print("python main.py --q-pos 10 --q-vel 1 --r-acc 1") 